#!/bin/bash
set -e

# 函数：识别操作系统类型和版本
identify_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS=$ID
        VERSION_ID=$VERSION_ID
        
        # 处理基于其他发行版的系统
        if [[ "$ID_LIKE" == *"debian"* ]]; then
            OS_FAMILY="debian"
        elif [[ "$ID_LIKE" == *"rhel"* || "$ID_LIKE" == *"fedora"* || "$ID_LIKE" == *"centos"* ]]; then
            OS_FAMILY="rhel"
        elif [[ "$ID" == *"suse"* ]]; then
            OS_FAMILY="suse"
        elif [[ "$ID" == *"arch"* ]]; then
            OS_FAMILY="arch"
        else
            OS_FAMILY=$OS
        fi
    elif [ -f /etc/redhat-release ]; then
        OS_FAMILY="rhel"
        OS=$(cat /etc/redhat-release | awk '{print tolower($1)}')
    elif [ -f /etc/debian_version ]; then
        OS_FAMILY="debian"
        OS="debian"
    else
        echo "无法识别的操作系统。"
        exit 1
    fi
    
    echo "检测到操作系统: $OS (家族: $OS_FAMILY)"
}

# 函数：获取主机IP
get_host_ip() {
    # 尝试多种方法获取主机IP，优先选择非环回地址
    HOST_IP=$(ip -4 addr show scope global | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | head -n 1)
    
    # 如果上面的方法失败，尝试其他方法
    if [ -z "$HOST_IP" ]; then
        HOST_IP=$(hostname -I | awk '{print $1}')
    fi
    
    # 如果仍然失败，再尝试其他方法
    if [ -z "$HOST_IP" ]; then
        HOST_IP=$(ifconfig | grep -Eo 'inet (addr:)?([0-9]*\.){3}[0-9]*' | grep -Eo '([0-9]*\.){3}[0-9]*' | grep -v '127.0.0.1' | head -n 1)
    fi
    
    # 如果所有方法都失败，使用localhost
    if [ -z "$HOST_IP" ]; then
        echo "警告: 无法确定主机IP，将使用127.0.0.1"
        HOST_IP="127.0.0.1"
    fi
    
    echo "使用主机IP: $HOST_IP"
}

# 函数：安装依赖工具
install_dependencies() {
    echo "正在安装依赖工具..."
    
    case $OS_FAMILY in
        debian)
            apt-get update -qq && apt-get install -y -qq curl wget net-tools tar iproute2 iputils-ping
            ;;
        rhel)
            if command -v dnf &> /dev/null; then
                dnf -y install curl wget net-tools tar iproute iputils
            else
                yum -y install curl wget net-tools tar iproute iputils
            fi
            ;;
        suse)
            zypper --non-interactive install curl wget net-tools tar iproute2 iputils
            ;;
        arch)
            pacman -Sy --noconfirm curl wget net-tools tar iproute2 iputils
            ;;
        *)
            echo "警告: 未知的包管理器，尝试使用通用方法安装依赖"
            if command -v apt-get &> /dev/null; then
                apt-get update -qq && apt-get install -y -qq curl wget net-tools tar iproute2 iputils-ping
            elif command -v dnf &> /dev/null; then
                dnf -y install curl wget net-tools tar iproute iputils
            elif command -v yum &> /dev/null; then
                yum -y install curl wget net-tools tar iproute iputils
            else
                echo "错误: 无法安装依赖，请手动安装curl、wget、net-tools、tar、iproute2和iputils"
                exit 1
            fi
            ;;
    esac
}

# 函数：设置系统服务和用户
setup_system() {
    echo "正在设置系统服务和用户..."
    
    # 创建vmagent配置文件目录
    mkdir -p /etc/victoriametrics/vmagent
    # 创建vmagent临时数据缓存目录
    mkdir -p /var/lib/vmagent-remotewrite-data

    # 检查victoriametrics组是否存在，不存在则创建
    if ! getent group victoriametrics > /dev/null 2>&1; then
        groupadd --system victoriametrics 2>/dev/null || true
    fi

    # 检查victoriametrics用户是否存在，不存在则创建
    if ! id -u victoriametrics > /dev/null 2>&1; then
        useradd --system --home-dir /var/lib/vmagent-remotewrite-data --no-create-home --shell /sbin/nologin \
            --gid victoriametrics victoriametrics 2>/dev/null || true
    fi

    chown -R victoriametrics:victoriametrics /var/lib/vmagent-remotewrite-data
}

# 函数：安装和设置服务
install_vmagent() {
    echo "正在安装vmagent..."
    
    # 获取vmagent最新版本
    VM_VERSION=$(curl -s "https://api.github.com/repos/VictoriaMetrics/VictoriaMetrics/tags" | grep '"name":' | grep -v 'beta\|alpha\|rc' | head -n 1 | awk -F '"' '{print $4}')
    echo "正在下载vmagent版本: $VM_VERSION"

    # 下载并安装vmagent
    TMP_DIR=$(mktemp -d)
    wget -q "https://github.com/VictoriaMetrics/VictoriaMetrics/releases/download/${VM_VERSION}/vmutils-linux-amd64-${VM_VERSION}.tar.gz" -O "${TMP_DIR}/vmutils.tar.gz"

    cd "${TMP_DIR}" && tar -xzf "${TMP_DIR}/vmutils.tar.gz" vmagent-prod
    mv "${TMP_DIR}/vmagent-prod" /usr/local/bin/
    chmod +x /usr/local/bin/vmagent-prod

    # 清理临时目录
    rm -rf "${TMP_DIR}"
}

# 函数：创建服务和配置文件
create_service_files() {
    echo "正在创建服务和配置文件..."
    
    # 创建systemd服务文件
    cat > /etc/systemd/system/vmagent.service <<EOF
[Unit]
Description=VictoriaMetrics Agent - scrapes and forwards metrics to VictoriaMetrics or any Prometheus-compatible storage
Documentation=https://docs.victoriametrics.com/vmagent.html
After=network.target

[Service]
Type=simple
User=victoriametrics
Group=victoriametrics
StartLimitBurst=5
StartLimitInterval=0
Restart=on-failure
RestartSec=3
EnvironmentFile=-/etc/victoriametrics/vmagent/vmagent.conf
ExecStart=/usr/local/bin/vmagent-prod \$ARGS
ExecStop=/bin/kill -s SIGTERM \$MAINPID
ExecReload=/bin/kill -HUP \$MAINPID
LimitNOFILE=1048576
LimitNPROC=1048576
LimitCORE=infinity
WorkingDirectory=/var/lib/vmagent-remotewrite-data
ReadWritePaths=/var/lib/vmagent-remotewrite-data
StandardOutput=journal
StandardError=journal
SyslogIdentifier=vmagent
PrivateTmp=yes
ProtectHome=yes
NoNewPrivileges=yes
ProtectSystem=strict
ProtectControlGroups=true
ProtectKernelModules=true
ProtectKernelTunables=yes

[Install]
WantedBy=multi-user.target
EOF

    # 创建配置文件
    cat > /etc/victoriametrics/vmagent/vmagent.conf <<EOF
ARGS="-promscrape.config=/etc/victoriametrics/vmagent/scrape.yml -remoteWrite.url=http://${HOST_IP}:8428/api/v1/write -remoteWrite.tmpDataPath=/var/lib/vmagent-remotewrite-data -promscrape.suppressScrapeErrors -envflag.enable"
EOF

    # 创建采集配置文件
    cat > /etc/victoriametrics/vmagent/scrape.yml <<EOF
global:
  scrape_interval: 10s
  scrape_timeout: 30s
  external_labels:
    host: "$(hostname)"
    ip: "${HOST_IP}"

scrape_configs:
  - job_name: 'vmagent'
    static_configs:
      - targets: ['${HOST_IP}:8429']
        labels:
          instance: "$(hostname)"
  - job_name: 'victoriametrics'
    static_configs:
      - targets: ['${HOST_IP}:8428']
        labels:
          instance: "$(hostname)"
EOF

    # 设置权限
    chown -R victoriametrics:victoriametrics /var/lib/vmagent-remotewrite-data
    chown -R victoriametrics:victoriametrics /etc/victoriametrics/vmagent
    chmod 644 /etc/systemd/system/vmagent.service
    chmod 644 /etc/victoriametrics/vmagent/vmagent.conf
    chmod 644 /etc/victoriametrics/vmagent/scrape.yml
}

# 函数：启用和启动服务
enable_service() {
    echo "正在启用和启动vmagent服务..."
    
    systemctl daemon-reload
    systemctl enable vmagent.service
    systemctl restart vmagent.service
    
    # 等待服务启动
    sleep 2
    
    if systemctl is-active --quiet vmagent.service; then
        echo "vmagent服务已成功启动"
    else
        echo "警告: vmagent服务可能未成功启动，请检查日志: journalctl -u vmagent.service"
    fi
}

# 函数：显示安装信息
show_info() {
    echo "=================== 安装信息 ==================="
    echo "vmagent版本: $VM_VERSION"
    echo "主机IP: $HOST_IP"
    echo "配置文件位置: /etc/victoriametrics/vmagent/"
    echo "数据缓存目录: /var/lib/vmagent-remotewrite-data"
    echo "服务状态:"
    systemctl status vmagent.service --no-pager | head -n 5
    echo "进程信息:"
    ps aux | grep vmagent-prod | grep -v grep || echo "未检测到vmagent进程"
    echo "==============================================="
    echo "vmagent安装和服务设置完成。"
    echo "可以通过以下命令查看详细日志: journalctl -u vmagent.service -f"
}

# 执行主要功能
main() {
    echo "开始安装VictoriaMetrics Agent..."
    
    # 检查是否以root身份运行
    if [ "$(id -u)" -ne 0 ]; then
        echo "错误: 此脚本需要root权限运行"
        exit 1
    fi
    
    identify_os
    install_dependencies
    get_host_ip
    setup_system
    install_vmagent
    create_service_files
    enable_service
    show_info
}

# 执行主函数
main